{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial for Robotic Modeling using Mujoco \n",
    "\n",
    "This tutorial covers the basics of the physics modeling with Mujoco. We will primarily focus on rigid-body systems that are prevalent in most robotic applications. \n",
    "\n",
    "## Principles of Physics Modeling\n",
    "In order to simulate a physical system, we typically need to tell the simulator three things:\n",
    "1. **System models**, often described as ODE/PDE or just functions. Often times, it is hard to derive analytically the system (differential) equations. Therefore, the most important functionality of a simulator is to allow the user to easily specify physical properties and automatically translate these properties into (differential) equations. For rigid body robotic systems, we need to tell the simulator:\n",
    "   - **Muti-body system**: articulated rigid body system consists of multiple rigid bodies connected through joints. We need to define each body and describe its physical properties (e.g. pos, orientation, inertia matrix, etc), and how it is connected to adjacent bodies (through joints) \n",
    "   - **Environment/interaction**: gravity direction, ground, table, obstacle, wall, person, other robots, etc. How the robot interacts with the environments depends on many physical properties that can be modified, including friction, contact, collision mechanisms, \n",
    "   - **Actuators/sensors**: robot's motion is completely determined by each joint's motion. Joint's motion is determined by sensing and actuation strategies. The simulation is about the closed-loop behavior of the robot's motion and its interaction with the environment. \n",
    "2. **Physics Engine**: This part determines how to solve the (differential) equations associated with the specified robot-environment physical system. Parameters that can be modified include simulation time step, integrator (Euler, RK4, implicit, implicitfast), tolerance (to terminate iterative solver), constraint solver (PGS, CG, Newton), number of iterations (for constraint solver), etc. For Mujoco xml, these parameters are mostly specified through the \\<option\\> element. \n",
    "3. **Visualization**: This part determines how to visualize the robotic system scene. In Mujoco xml, some key visualization components are listed below:\n",
    "   -  **light**: body/light will create a light, which moves with the body where it is defined. To create a fixed light, define it in the world body. The lights created here are in addition to the default headlight which is always defined and is adjusted via the visual element. The maximum number of lights that can be active simultaneously is 8, counting the headlight.\n",
    "   -  **Camera**: body/⁠camera creates a camera, which moves with the body where it is defined. To create a fixed camera, define it in the world body. ca\n",
    "   -  **geom**: body/⁠geom creates a geom, and attaches it rigidly to the body within which the geom is defined. Multiple geoms can be attached to the same body. At runtime they determine the appearance and collision properties of the body. \n",
    "   -  **Material**: asset/⁠material creates a material asset. It can be referenced from skins, geoms, sites and tendons to set their appearance. For rigid-body system, material is referenced from geom associated with a body. \n",
    "   - **Body vs Geom**: \\<body\\>: Represents a rigid entity with mass, inertia, and movable via joints. It does not have a shape until \\<geom\\> elements are attached to it. While \\<geom\\>: Defines the shape and size, used for collisions and visualization. It does not have mass or inertia, but can have density specified to infer mass. \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Python Examples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 1: Geometric objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import mujoco\n",
    "import mujoco.viewer as viewer\n",
    "\n",
    "xmlModel = \"\"\"\n",
    "<mujoco>\n",
    "  <worldbody>\n",
    "    <geom name=\"red_box\" type=\"box\" pos=\"0 0 0\" size=\".2 .2 .2\" rgba=\"1 0 0 1\"/>\n",
    "    <geom name=\"blue_sphere\" pos=\".2 .2 .2\" size=\".1\" rgba=\"0 0 1 1\"/>\n",
    "  </worldbody>\n",
    "</mujoco>\n",
    "\"\"\"\n",
    "model = mujoco.MjModel.from_xml_string(xmlModel)\n",
    "data = mujoco.MjData(model)\n",
    "viewer.launch(model,data)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The box and sphere do not move. They are geometric objects only for visulization. In order for them to move, they need to be attached to a body with motion freedom specified by a joint. Before adding joint, we first use asset to specify material and texture, use \\<default\\> to change the default setting for the geom element"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "xml_ground = \"\"\"\n",
    "<mujoco>\n",
    "  <asset>  \n",
    "    <texture name=\"grid\" type=\"2d\" builtin=\"checker\" rgb1=\".1 .2 .3\"\n",
    "     rgb2=\".2 .3 .4\" width=\"300\" height=\"300\" mark=\"none\"/>\n",
    "    <material name=\"grid\" texture=\"grid\" texrepeat=\"6 6\" texuniform=\"true\" reflectance=\".2\"/>\n",
    "    <material name=\"wall\" rgba='.5 .5 .5 1'/>\n",
    "  </asset>\n",
    "  <default>\n",
    "    <geom type=\"box\" size=\".05 .05 .05\" />\n",
    "  </default>\n",
    "\n",
    "  <worldbody>\n",
    "    <light name=\"light\" pos=\"-.2 0 1\"/>\n",
    "    <geom name=\"ground\" type=\"plane\" size=\".5 .5 10\" material=\"grid\" friction=\".1\"/>\n",
    "    <camera name=\"y\" pos=\"-.1 -.6 .3\" xyaxes=\"1 0 0 0 1 2\"/>\n",
    "    <geom name=\"red_box\" pos=\"0 0 0.2\" type=\"box\" size=\".2 .2 .2\" rgba=\"1 0 0 1\"/>\n",
    "    <geom name=\"blue_sphere\" pos=\".2 .2 .2\" size=\".1\" rgba=\"0 0 1 1\"/>\n",
    "  </worldbody>\n",
    "</mujoco>\n",
    "\"\"\"\n",
    "model = mujoco.MjModel.from_xml_string(xml_ground)\n",
    "data = mujoco.MjData(model)\n",
    "viewer.launch(model,data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 2： Multi-body Robotic Systems\n",
    "\n",
    "Now let's add a body and a joint to the world frame (worldbody). Simulate the physical motion caused by gravity. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "xml = \"\"\"\n",
    "<mujoco>\n",
    "  <asset>  \n",
    "    <texture name=\"grid\" type=\"2d\" builtin=\"checker\" rgb1=\".1 .2 .3\"\n",
    "     rgb2=\".2 .3 .4\" width=\"300\" height=\"300\" mark=\"none\"/>\n",
    "    <material name=\"grid\" texture=\"grid\" texrepeat=\"6 6\" texuniform=\"true\" reflectance=\".2\"/>\n",
    "    <material name=\"wall\" rgba='.5 .5 .5 1'/>\n",
    "  </asset>\n",
    "  <option gravity=\"0 0 -9.81\"/>        \n",
    "  <worldbody>\n",
    "  <geom name=\"ground\" type=\"plane\" pos=\"0 0 -1\" size=\"10 10 10\" material=\"grid\" friction=\".1\"/>\n",
    "    <light name=\"top\" pos=\"0 0 1\"/>\n",
    "    <body name=\"box_and_sphere\" euler=\"0 0 -30\">\n",
    "      <joint name=\"swing\" type=\"hinge\" axis=\"0 1 0\" pos=\"-.2 -.2 -.2\"/>\n",
    "      \n",
    "      <geom name=\"red_box\" type=\"box\" size=\".2 .2 .2\" rgba=\"1 0 0 1\"/>\n",
    "      <geom name=\"green_sphere\" pos=\".2 .2 .2\" size=\".1\" rgba=\"0 1 0 1\"/>\n",
    "    </body>\n",
    "  </worldbody>\n",
    "</mujoco>\n",
    "\"\"\"\n",
    "model = mujoco.MjModel.from_xml_string(xml)\n",
    "data = mujoco.MjData(model)\n",
    "viewer.launch(model,data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's have a scene with ground (modeled as plane) and two boxes. The box can move freely in space, which is modeled as a \"free\" joint relative to the worldbody. You can modify the friction coefficients to see the resulting sliding motion. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "xml3 = \"\"\"\n",
    "<mujoco>\n",
    "  <asset>  \n",
    "    <texture name=\"grid\" type=\"2d\" builtin=\"checker\" rgb1=\".1 .2 .3\"\n",
    "     rgb2=\".2 .3 .4\" width=\"300\" height=\"300\" mark=\"none\"/>\n",
    "    <material name=\"grid\" texture=\"grid\" texrepeat=\"6 6\"\n",
    "     texuniform=\"true\" reflectance=\".2\"/>\n",
    "     <material name=\"wall\" rgba='.5 .5 .5 1'/>\n",
    "  </asset>\n",
    "  <default>\n",
    "    <geom type=\"box\" size=\".05 .05 .05\" />\n",
    "    <joint type=\"free\"/>\n",
    "  </default>\n",
    "\n",
    "  <worldbody>\n",
    "    <light name=\"light\" pos=\"-.2 0 1\"/>\n",
    "    <geom name=\"ground\" type=\"plane\" size=\"10 10 10\" material=\"grid\"\n",
    "     zaxis=\"-.2 0 1\" friction=\".2\"/>\n",
    "    <camera name=\"y\" pos=\"-.1 -.6 .3\" xyaxes=\"1 0 0 0 1 2\"/>\n",
    "    <body pos=\"0 1 .3\">\n",
    "      <joint/>\n",
    "      <geom friction=\"0.2\"/>\n",
    "    </body>\n",
    "    <body pos=\"0 0 .3\">\n",
    "      <joint/>\n",
    "      <geom friction=\"0.4\"/>\n",
    "    </body>\n",
    "  </worldbody>\n",
    "\n",
    "</mujoco>\n",
    "\"\"\"\n",
    "\n",
    "# load\n",
    "model = mujoco.MjModel.from_xml_string(xml3)\n",
    "data = mujoco.MjData(model)\n",
    "viewer.launch(model,data)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's have a multi-body robot\n",
    "\n",
    "To construct a multi-body robot, one needs to clearly specify\n",
    "- body frame\n",
    "- geom frame (the frame of geomoetric object)\n",
    "- joint frame\n",
    "For simplicity, we shall try to let the inertia frame coincide with the body frame. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import mujoco\n",
    "import mujoco.viewer as viewer\n",
    "\n",
    "xml4=\"\"\"<mujoco model=\"3R_robot\">\n",
    "    <compiler angle=\"degree\"/>\n",
    "    <asset>  \n",
    "        <texture name=\"grid\" type=\"2d\" builtin=\"checker\" rgb1=\".1 .2 .3\"\n",
    "        rgb2=\".2 .3 .4\" width=\"300\" height=\"300\" mark=\"none\"/>\n",
    "        <material name=\"grid\" texture=\"grid\" texrepeat=\"6 6\" texuniform=\"true\" reflectance=\".2\"/>\n",
    "    </asset>\n",
    "    <default>\n",
    "        <joint type=\"hinge\" axis=\"0 0 1\" limited=\"true\"/>\n",
    "        <geom type=\"cylinder\" size=\".025 .1\" />\n",
    "    </default>\n",
    "\n",
    "    <worldbody>\n",
    "        <light diffuse=\".5 .5 .5\" pos=\"0 0 3\" dir=\"0 0 -1\"/>\n",
    "        <geom type=\"plane\" size=\"1 1 0.1\" material=\"grid\"/>        \n",
    "        <body name=\"BaseLink\" pos=\"0 0 .1\">  \n",
    "            <geom type=\"cylinder\" pos=\"0 0 0\" size=\".025 .1\" />  \n",
    "            <body name=\"link1\" pos=\"0 0.1 0.125\" euler=\"-90 0 0\">\n",
    "                <joint name=\"joint1\" pos=\"0 0 -0.1\" range=\"-90 90\" axis =\"0 1 0\"/>\n",
    "                <geom pos=\"0 0 0\" rgba=\".6 .2 .2 1\"/>\n",
    "                <site name=\"torque_site\" pos=\"0 0.2 0\"/>\n",
    "                <body name=\"link2\" pos=\"0 0 0.2\">\n",
    "                    <joint name=\"joint2\" pos=\"0 0 -0.1\" range=\"-90 90\" axis=\"0 1 0\"/>\n",
    "                    <geom rgba=\".2 .6 1 1\"/>\n",
    "                    <site name=\"end_effector\" pos=\"0 0 0.1\" size=\"0.01\"/>\n",
    "                </body>\n",
    "            </body>\n",
    "        </body>\n",
    "    </worldbody>\n",
    "</mujoco>\"\"\"\n",
    "# load\n",
    "model = mujoco.MjModel.from_xml_string(xml4)\n",
    "data = mujoco.MjData(model)\n",
    "mujoco.mj_resetDataKeyframe(model, data, 0)\n",
    "viewer.launch(model,data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example 3: Multi-body robot with control\n",
    "We add actuators to the joints. The first joint is a torque controlled joint while the second joint is a position-joint. We set the inital pose to be $(0,\\frac{\\pi}{2})$ while the control inputs are zeros by default. So the simulation will show the process of the second joint rotating towards 0 degree and eventually settles down at zero. You can modify the damping kv term to see the response. Since we did not have any control over the first joint (i.e. desired torque is always zero), the first joint may rotate due to the second body's motion. We can add damping to the first joint. This can be done directly through the \\<joint\\/> element. \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import mujoco\n",
    "import mujoco.viewer as viewer\n",
    "\n",
    "xml4=\"\"\"<mujoco model=\"3R_robot\">\n",
    "    <compiler angle=\"degree\"/>\n",
    "    <asset>  \n",
    "        <texture name=\"grid\" type=\"2d\" builtin=\"checker\" rgb1=\".1 .2 .3\"\n",
    "        rgb2=\".2 .3 .4\" width=\"300\" height=\"300\" mark=\"none\"/>\n",
    "        <material name=\"grid\" texture=\"grid\" texrepeat=\"6 6\" texuniform=\"true\" reflectance=\".2\"/>\n",
    "    </asset>\n",
    "    <default>\n",
    "        <joint type=\"hinge\" axis=\"0 0 1\" limited=\"true\"/>\n",
    "        <geom type=\"cylinder\" size=\".025 .1\" />\n",
    "    </default>\n",
    "\n",
    "    <worldbody>\n",
    "        <light diffuse=\".5 .5 .5\" pos=\"0 0 3\" dir=\"0 0 -1\"/>\n",
    "        <geom type=\"plane\" size=\"1 1 0.1\" material=\"grid\"/>        \n",
    "        <body name=\"BaseLink\" pos=\"0 0 0.1\">  \n",
    "            <geom type=\"cylinder\" pos=\"0 0 0\" size=\".025 .1\" />  \n",
    "            <body name=\"link1\" pos=\"0 0.1 0.125\" euler=\"-90 0 0\">\n",
    "                <joint name=\"joint1\" pos=\"0 0 -0.1\" range=\"-90 90\" axis =\"0 1 0\" damping=\"0.1\"/>\n",
    "                <geom pos=\"0 0 0\" rgba=\".6 .2 .2 1\"/>\n",
    "                <body name=\"link2\" pos=\"0 0 0.2\">\n",
    "                    <joint name=\"joint2\" pos=\"0 0 -0.1\" range=\"-90 90\" axis=\"0 1 0\"/>\n",
    "                    <geom rgba=\".2 .6 1 1\"/>\n",
    "                    <site name=\"end_effector\" pos=\"0 0 0.1\" size=\"0.01\"/>\n",
    "                </body>\n",
    "            </body>\n",
    "        </body>\n",
    "    </worldbody>\n",
    "    <!-- Actuators for controlling the joints -->\n",
    "    <actuator>\n",
    "        <motor name=\"torque_control\" joint=\"joint1\" ctrlrange=\"-10 10\"/> \n",
    "        <position name=\"position_control\" joint=\"joint2\" kp=\"2\" kv=\"0.02\" ctrlrange=\"-90 90\"/>\n",
    "    </actuator>\n",
    "    <!-- Sensors for monitoring joint positions and torques -->\n",
    "    <sensor>\n",
    "        <jointpos name=\"joint1_position_sensor\" joint=\"joint1\"/> <!-- Position sensor for joint1 -->\n",
    "        <jointpos name=\"joint2_position_sensor\" joint=\"joint2\"/> <!-- Position sensor for joint2 -->\n",
    "    </sensor>\n",
    "    <keyframe>\n",
    "        <!-- Initial configuration -->\n",
    "        <key name=\"initial_pose\" qpos=\"0 -1.57\"/>        \n",
    "        <!-- Bent configuration -->\n",
    "        <key name=\"bent_pose\" qpos=\"1.57079 0.78539\"/>\n",
    "    </keyframe>\n",
    "\n",
    "</mujoco>\"\"\"\n",
    "\n",
    "model = mujoco.MjModel.from_xml_string(xml4)\n",
    "data = mujoco.MjData(model)\n",
    "mujoco.mj_resetDataKeyframe(model, data, 1)\n",
    "viewer.launch(model,data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "#height map example\n",
    "\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "import time\n",
    "import mujoco\n",
    "import mujoco.viewer\n",
    "\n",
    "height_map = np.random.rand(100, 100) * 255  \n",
    "height_map = height_map.astype(np.uint8)\n",
    "\n",
    "Image.fromarray(height_map).save(\"terrain.png\")\n",
    "\n",
    "\n",
    "xml_string = \"\"\"\n",
    "<mujoco model=\"height_field_terrain\">\n",
    "    <asset>\n",
    "        <hfield name=\"terrain_heightfield\" file=\"terrain.png\" size=\"50 50 1 1\" />\n",
    "    </asset>\n",
    "\n",
    "    <worldbody>\n",
    "        <light directional=\"true\" diffuse=\".8 .8 .8\" pos=\"0 0 10\" dir=\"0 0 -10\"/>\n",
    "        <geom name=\"terrain_hfield\" type=\"hfield\" hfield=\"terrain_heightfield\" pos=\"0 0 0\" size=\"5 5 1\"/>\n",
    "        <geom name=\"terrain_box\" type=\"box\" pos=\"0 0 2\" size=\"1 1 1\" rgba=\"1 0 0 1\"/>\n",
    "        <geom name=\"ground\" type=\"plane\" pos=\"5 0 3\" size=\"2 2 0.1\" rgba=\"0 0 1 1\"/>\n",
    "    </worldbody>\n",
    "</mujoco>\n",
    "\"\"\"\n",
    "\n",
    "m = mujoco.MjModel.from_xml_string(xml_string)\n",
    "d = mujoco.MjData(m)\n",
    "viewer.launch(m,d)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sim",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
